package io.summer.gateway.filter;

import static java.util.function.Function.identity;

import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import io.summer.gateway.tool.MediaTypeTool;
import org.reactivestreams.Publisher;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.cloud.gateway.filter.factory.rewrite.GzipMessageBodyResolver;
import org.springframework.cloud.gateway.filter.factory.rewrite.MessageBodyDecoder;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;

import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;

@Slf4j
class LogHttpResponseDecorator extends ServerHttpResponseDecorator {
  private final StringBuilder logCollector;
  private final Long startTime;
  private final String separator;
  private final Map<String, MessageBodyDecoder> messageBodyDecoders;
  private static final int MAX_DECODE_CONTENT_SIZE = 3 * 1024;

  public LogHttpResponseDecorator(final ServerHttpResponse delegate, final StringBuilder logCollector, final String sep) {
    super(delegate);
    this.logCollector = logCollector;
    this.startTime = System.currentTimeMillis();
    this.separator = sep;
    this.messageBodyDecoders = Stream.of(new GzipMessageBodyResolver())
      .collect(Collectors.toMap(MessageBodyDecoder::encodingType, identity()));
  }

  private void recordLog(final String response, int rawLen) {
    if (logCollector == null) {
      return;
    }

    logCollector.append(separator);
    logCollector.append("<<<<<<<<<<<<<<<< Response").append(separator);
    if (startTime != null) {
      logCollector.append("Cost: ").append(System.currentTimeMillis() - startTime).append("ms");
      logCollector.append(separator);
    }
    logCollector.append("Status: ").append(this.getStatusCode()).append(separator);
    logCollector.append("Header: ").append(this.getHeaders()).append(separator);
    logCollector.append("Response-Length: ").append(rawLen).append(separator);
    logCollector.append("Response: ").append(response).append(separator);
    if (log.isInfoEnabled()) {
      log.info(logCollector.toString());
    }
  }

  private boolean exceedDecodeThrottle(byte[] contents) {
    if (contents == null) {
      return false;
    }
    return contents.length > MAX_DECODE_CONTENT_SIZE;
  }

  @Override
  public Mono<Void> writeWith(final Publisher<? extends DataBuffer> body) {
    if (body instanceof Flux) {
      final Flux<? extends DataBuffer> fluxBody = (Flux<? extends DataBuffer>) body;

      Mono<DataBuffer> mono = DataBufferUtils.join(fluxBody)
        .map(dataBuffer -> {
          byte[] buffer = new byte[dataBuffer.readableByteCount()];
          dataBuffer.read(buffer);
          DataBufferUtils.release(dataBuffer);
          return buffer;
        })
        .publishOn(Schedulers.parallel())
        .flatMap(bytes -> {
            decodeLog(bytes);
            return Mono.just(getDelegate().bufferFactory().wrap(bytes));
          }
        );
      return super.writeWith(mono);
    }
    return super.writeWith(body);
  }

  private void decodeLog(byte[] bytes) {
    byte[] contents = bytes;
    boolean exceedDecodeThrottle = exceedDecodeThrottle(contents);
    if (exceedDecodeThrottle) {
      recordLog(String.format("response body exceed (%d)bytes, no printing for it", MAX_DECODE_CONTENT_SIZE), contents.length);
      return;
    }

    List<String> encodingHeaders = getHeaders()
      .getOrEmpty(HttpHeaders.CONTENT_ENCODING);
    if (!encodingHeaders.isEmpty()) {
      String decoderName = encodingHeaders.get(0);
      MessageBodyDecoder decoder = messageBodyDecoders.get(decoderName);
      if (decoder != null) {
        try {
          contents = decoder.decode(bytes);
        } catch (Exception e) {
          log.error("log-filter decode data of type {} failed.", decoder.encodingType());
        }
      }
    }
    recordLog(new String(contents, StandardCharsets.UTF_8), Objects.isNull(bytes) ? 0 : bytes.length);
  }
}

/**
 * @author Moonlight
 */
@Component
@ConditionalOnProperty(prefix = "summer.gateway.log", value = "enable", havingValue = "true")
public class LogGlobalFilter implements GlobalFilter, Ordered {
  private static final String CACHE_REQUEST_BODY_OBJECT_KEY = "cachedRequestBodyObject";
  @Value("${summer.gateway.log.separator:\n}")
  private String separator;

  @Override
  public Mono<Void> filter(final ServerWebExchange exchange, final GatewayFilterChain chain) {
    final StringBuilder logCollector = logCollector(exchange.getRequest());
    recordLog(exchange, logCollector);
    final LogHttpResponseDecorator serverHttpResponse = new LogHttpResponseDecorator(exchange.getResponse(),
      logCollector, separator);
    return chain.filter(exchange.mutate().response(serverHttpResponse).build());
  }

  private StringBuilder logCollector(final ServerHttpRequest request) {
    final String uri = request.getURI().getRawPath();
//        if (uri.contains(MySwaggerResourcesProvider.DOCS_URI)) {
//          return null;
//        }
    return new StringBuilder(1024);
  }

  private void recordLog(final ServerWebExchange exchange, final StringBuilder log) {
    if (log == null) {
      return;
    }

    final ServerHttpRequest request = exchange.getRequest();
    log.append("Request >>>>>>>>>>>>>>>>>").append(separator);
    log.append("Url: ").append(request.getURI().getRawPath()).append(separator);
    log.append("Method: ").append(request.getMethodValue()).append(separator);
    log.append("Query: ").append(request.getQueryParams()).append(separator);
    final HttpHeaders headers = request.getHeaders();
    log.append("Headers: ").append(headers).append(separator);
    final InetSocketAddress address = request.getRemoteAddress();
    if (address != null) {
      log.append("Address: ").append(address.getHostString()).append(":").append(address.getPort()).append(separator);
    }
    if (!MediaTypeTool.fileType(headers.getContentType())) {
      final String requestBody = exchange.getAttribute(CACHE_REQUEST_BODY_OBJECT_KEY);
      log.append("Body: ").append(requestBody);
    } else {
      log.append("Body: ").append("omit printing body file, len: ").append(headers.getContentLength());
    }
    log.append(separator);
  }


  @Override
  public int getOrder() {
    return Ordered.HIGHEST_PRECEDENCE;
  }
}
